from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from scipy.spatial import KDTree
from sklearn.linear_model import LinearRegression
from scipy.stats import norm

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

model_forward_vmap = jax.vmap(model_forward, in_axes=(None, 0, None, None))

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

# 全局变量，用于存储和图像显示线程交互的数据
class imgview:
    global_image = None
    imgview_exit = False
    trajectory = []
    focus_i = 0
    traj_i = 0

imgview_data = imgview()

# 定义一个函数，用于在独立线程中显示图像
def show_image():
    grid_size_display = 20
    while not imgview_data.imgview_exit:
        # 检查全局变量是否有图像
        if imgview_data.global_image is not None:
            img = np.copy(imgview_data.global_image)
            state_x = imgview_data.trajectory[imgview_data.traj_i][0]
            state_y = imgview_data.trajectory[imgview_data.traj_i][1]
            cv2.circle(img, (state_y * grid_size_display + int(grid_size_display/2), state_x * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
            # 显示图像
            cv2.imshow("Image", img)
            key = cv2.waitKey(1)
            if key == ord('a'):
                imgview_data.focus_i -= 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
            elif key == ord('d'):
                imgview_data.focus_i += 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
        else:
            # 图像还未产生，等待100毫秒
            time.sleep(0.1)

# 计算直方图并且显示
def view_histogram(data):

    strength_hist, bins = np.histogram(data, bins = 100)
    strength_hist = np.array(strength_hist)
    bin_centers = (bins[:-1] + bins[1:]) / 2
    # 对 strength_hist 可视化
    plt.figure()
    plt.plot(bin_centers, strength_hist)
    plt.show()

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    """ 1. 计算所有随机位置的向量场
    """

    n_samples = 4000000
    # 创建一批随机的128维向量，其每个维度分布在[-1, 1]之间，共计n_samples个
    random_states = np.random.uniform(-1, 1, (n_samples, 128))
    obs_zero = jnp.zeros((10,))

    # # 先自由运行一段时间，使其到达慢区
    # for i in range(100):
    #     progress_bar(i, 100)
    #     random_states, _ = model_forward_vmap(params, random_states, obs_zero, model)

    states_vector, _ = model_forward_vmap(params, random_states, obs_zero, model)
    states_vector = states_vector - random_states

    print("shape of random_states: ", random_states.shape)
    print("shape of state_vectors: ", states_vector.shape)

    # 显示这些变量的内存大小，以GiB为单位
    print("size of random_states: ", random_states.nbytes / 1024 / 1024 / 1024)
    print("size of state_vectors: ", states_vector.nbytes / 1024 / 1024 / 1024)

    """ 2. 进行数据处理, 提取慢点、LDS等
    """
    
    keep_ratio = 0.01
    keep_thresh = 4.5
    keep_n_samples = 2000
    
    vector_strength = jnp.linalg.norm(states_vector, axis = 1)
    print("shape of vector_strength: ", vector_strength.shape)
    # view_histogram(vector_strength)

    # # 保留 vector_strength 小于 keep_thresh 的部分作为慢点
    # slow_states = random_states[vector_strength < keep_thresh]
    # slow_states_vector = states_vector[vector_strength < keep_thresh]
    # print("shape of slow_states: ", slow_states.shape)
    # print("shape of slow_states_vector: ", slow_states_vector.shape)

    # 保留 vector_strength 最小的 keep_n_samples 个样本作为慢点
    idx = vector_strength.argsort()[:keep_n_samples]
    slow_states = random_states[idx]
    slow_states_vector = states_vector[idx]
    print("shape of slow_states: ", slow_states.shape)
    print("shape of slow_states_vector: ", slow_states_vector.shape)

    # 将 slow_states 进行 PCA 并且可视化前三维
    pca = PCA()
    pca.fit(slow_states)

    # 打印 variance ratio
    print("variance ratio: ", pca.explained_variance_ratio_)

    slow_states_pca = pca.transform(slow_states)
    print("shape of slow_states_pca: ", slow_states_pca.shape)
    # 对 slow_states_pca 可视化
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    ax.scatter(slow_states_pca[:,0], slow_states_pca[:,1], slow_states_pca[:,2], s = 0.1)
    plt.show()
    


if __name__ == "__main__":
    
    ivf()